import torch
import torch.nn as nn
# from networks.encoder import ResNet
# from networks.bridge_network_cell import TF_model_cell
from .bridge_network_nerve import TF_model_nerve
# from networks.decoder_baseline import Decoder
from .decoder_cell_k import Decoder_K
# from networks.decoder_cell import Decoder
from torchvision import models
# from networks.decoder_nerve import Decoder_1



class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace= True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
    def forward(self, input):
        return self.conv(input)




class Encoder(nn.Module):
    def __init__(self, in_channels=1):
        super(Encoder, self).__init__()
        # self.encoder = ResNet()

        model = models.resnet34(pretrained=True)
        self.conv_input = nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=7, stride=2, padding=3)
        self.bn1 = model.bn1
        self.relu = model.relu
        self.maxpool = model.maxpool

        self.layer1 = model.layer1
        self.layer2 = model.layer2
        self.layer3 = model.layer3
        self.layer4 = model.layer4


        self.model_cell2 = TF_model_nerve(d_model=240,
                                         nhead=8,
                                         d_ffn=1024,
                                         dropout=0.1,
                                         act="gelu",
                                         n_points=4,
                                         n_levels=1,
                                         n_sa_layers=1,
                                         in_channles=[64, 64, 128, 256, 512],
                                         proj_idxs=(2,),
                                         activation="relu")
        self.model_cell3 = TF_model_nerve(d_model=240,
                                          nhead=8,
                                          d_ffn=1024,
                                          dropout=0.1,
                                          act="gelu",
                                          n_points=4,
                                          n_levels=1,
                                          n_sa_layers=1,
                                          in_channles=[64, 64, 128, 256, 512],
                                          proj_idxs=(3,),
                                          activation="relu")
        self.model_cell4 = TF_model_nerve(d_model=240,
                                          nhead=8,
                                          d_ffn=1024,
                                          dropout=0.1,
                                          act="gelu",
                                          n_points=4,
                                          n_levels=1,
                                          n_sa_layers=1,
                                          in_channles=[64, 64, 128, 256, 512],
                                          proj_idxs=(4,),
                                          activation="relu")
        # self.model_nerve = TF_model_nerve()
        self.decoder = Decoder_K()
        self.pool = nn.MaxPool2d(2)

        # self.conv_11 = nn.ConvTranspose2d(368, 128, 2, stride=2)
        self.conv_11 = DoubleConv(368, 128)
        self.conv_22 = DoubleConv(128, 256)
        self.conv_33 = DoubleConv(496, 256)
        self.conv_44 = DoubleConv(752, 512)


    def forward(self, input):
        # feats = self.encoder(input)

        # 0 : torch.Size([1, 64, 192, 192])
        # 1 : torch.Size([1, 64, 96, 96])
        # 2 : torch.Size([1, 128, 48, 48])
        # 3 : torch.Size([1, 256, 24, 24])
        # 4 : torch.Size([1, 512, 12, 12])
        out = []
        y0 = self.conv_input(input)
        out.append(y0)
        y1 = self.layer1(self.maxpool(self.relu(self.bn1(y0))))
        out.append(y1)

        y2 = self.layer2(y1)
        out.append(y2)      ## 前三层
        trans_2 = self.model_cell2(out)

        out[-1] = torch.cat((out[-1], trans_2[-1]), dim=1)
        out[-1] = self.conv_11(out[-1])
        res = self.layer3(out[-1])
        out.append(res)

        trans_3 = self.model_cell3(out)

        out[-1] = torch.cat((out[-1], trans_3[-1]), dim=1)
        out[-1] = self.conv_33(out[-1])
        res_new = self.layer4(out[-1])
        out.append(res_new)

        trans_4 = self.model_cell4(out)
        out[-1] = torch.cat((out[-1], trans_4[-1]), dim=1)
        out[-1] = self.conv_44(out[-1])

        return  out


# out.shape torch.Size([1, 64, 192, 192])
# out.shape torch.Size([1, 64, 96, 96])
# out.shape torch.Size([1, 128, 48, 48])
# out.shape torch.Size([1, 256, 24, 24])
# out.shape torch.Size([1, 512, 12, 12])










if __name__ == '__main__':
    img = torch.Tensor(1, 1, 384, 384).cuda()
    net = NF_K().cuda()
    out = net(img)  ##
    # print("out.shape",out.shape) ###torch.Size([8, 2, 384, 384])
    # out_n = out[:,1,:,:]   ## 冒号取代像素值  torch.Size([8, 384, 384])
    # print(out_n.shape)
